package nl.cypherpunk.proLearner;

import de.learnlib.algorithms.lstar.mealy.ExtensibleLStarMealyBuilder;
import de.learnlib.api.algorithm.LearningAlgorithm;
import de.learnlib.api.logging.LearnLogger;
import de.learnlib.api.oracle.EquivalenceOracle;
import de.learnlib.api.query.DefaultQuery;
import de.learnlib.filter.cache.mealy.MealyCacheOracle;
import de.learnlib.filter.cache.mealy.MealyCaches;
import de.learnlib.filter.statistic.Counter;
import de.learnlib.filter.statistic.oracle.MealyCounterOracle;
import de.learnlib.oracle.membership.SULOracle;
import de.learnlib.util.statistics.SimpleProfiler;
import net.automatalib.automata.transducers.MealyMachine;
import net.automatalib.serialization.dot.GraphDOT;
import net.automatalib.words.Alphabet;
import net.automatalib.words.Word;
import nl.cypherpunk.proLearner.fuzzing.SUT;
import nl.cypherpunk.proLearner.mealyDot.MealyDot;
import nl.cypherpunk.proLearner.mealyDot.MealyState;
import nl.cypherpunk.proLearner.mealyDot.OutputTransition;
import nl.cypherpunk.proLearner.rtsp.RTSPConfig;
import nl.cypherpunk.proLearner.rtsp.RTSPSUL;
import nl.cypherpunk.proLearner.smtp.SMTPConfig;
import nl.cypherpunk.proLearner.smtp.SMTPSUL;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.text.DecimalFormat;
import java.util.*;

import static nl.cypherpunk.proLearner.Common.*;

public class ProLearner {
    LearningConfig config;
    Alphabet<String> alphabet;
    Map<String,Integer> subtypes;
    LinkedHashMap<String,List<Integer>> addtypes;
    boolean combine_query = false;
    SUT<String, String> sul;
    SULOracle<String, String> memOracle;
    ModifiedLogOracle.MealyLogOracle<String, String> logMemOracle;
    MealyCounterOracle<String, String> statsMemOracle;
    MealyCacheOracle<String, String> cachedMemOracle;
    MealyCounterOracle<String, String> statsCachedMemOracle;
    LearningAlgorithm<MealyMachine<?, String, ?, String>, String, Word<String>> learningAlgorithm;

    SULOracle<String, String> eqOracle;
    ModifiedLogOracle.MealyLogOracle<String, String> logEqOracle;
    MealyCounterOracle<String, String> statsEqOracle;
    MealyCacheOracle<String, String> cachedEqOracle;
    MealyCounterOracle<String, String> statsCachedEqOracle;
    EquivalenceOracle<MealyMachine<?, String, ?, String>, String, Word<String>> equivalenceAlgorithm;

    ArrayList<List<byte[]>> seedPool;
    ArrayList<List<String>> inputTypeSeqs;
    ArrayList<List<String>> outputTypeSeqs;
    ArrayList<DefaultQuery<String, Word<String>>> CEList;
    Boolean dedupFlag;
    public ProLearner(LearningConfig config) throws Exception {
        this.config = config;

        // Create output directory if it doesn't exist
        Path path = Paths.get(config.output_dir);
        if(Files.notExists(path)) {
            Files.createDirectories(path);
        }

        LearnLogger log = LearnLogger.getLogger(Learner.class.getSimpleName());

        // Check the type of learning we want to do and create corresponding configuration and SUL

        if(config.type == LearningConfig.TYPE_SMTP) {
            log.logPhase("Using SMTP SUL");
            String killall = "sudo killall exim";
            ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
            Process tmp = pbkill.start();
            tmp.waitFor();
            killall = "sudo rm /home/exim*";
            pbkill = new ProcessBuilder(killall.split(" "));
            tmp = pbkill.start();
            tmp.waitFor();
            // Create the SMTP SUL
            sul = new SMTPSUL(new SMTPConfig(config));
            alphabet = ((SMTPSUL)sul).getAlphabet();
        }
        else if(config.type == LearningConfig.TYPE_RTSP) {
            log.logPhase("Using RTSP SUL");
            String killall = "sudo killall testOnDemandRTSPServer";
            ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
            Process tmp = pbkill.start();
            tmp.waitFor();
            // Create the SMTP SUL
            sul = new RTSPSUL(new RTSPConfig(config));
            alphabet = ((RTSPSUL)sul).getAlphabet();
        }
    }


    public void learn() throws IOException, InterruptedException {
        /*try{
            test();
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }*/
        LearnLogger log = LearnLogger.getLogger(Learner.class.getSimpleName());
        initMapper();
        try{
            deterMutate();
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }
        int i = 1;
        addtypes = null;
        Counter progress = new Counter("Progress rounds", "");
        CEList = new ArrayList<>();
        dedupFlag = false;
        while(true) {
            progress.increment();
            System.out.println("Starting progress: " + progress.getCount());
            System.out.println("The index of seed: "+i);
            if(addtypes==null){
                logEqOracle = new ModifiedLogOracle.MealyLogOracle<String, String>(this.sul, LearnLogger.getLogger("queries"));
                statsEqOracle = new MealyCounterOracle<String, String>(logEqOracle, "MembershipQuery");
                cachedMemOracle = MealyCaches.createTreeCache(this.alphabet, statsEqOracle);
            }
            else {
                if(dedupFlag)
                {
                    cachedMemOracle = MealyCaches.createTreeCache(this.alphabet, statsEqOracle);
                }
                else {
                    Iterator iter = addtypes.entrySet().iterator();
                    while (iter.hasNext()) {
                        Map.Entry entry = (Map.Entry) iter.next();
                        String key = (String) entry.getKey();
                        List<Integer> val = (List<Integer>) entry.getValue();
                        for (int index = 0; index < val.size(); index++) {
                            String tmp = key + (new DecimalFormat("00").format(val.get(index)));
                            cachedMemOracle.addAlphabetSymbol(tmp);
                        }
                    }
                }
            }
            //dedupFlag = false;
            statsCachedMemOracle = new MealyCounterOracle<String, String>(cachedMemOracle, "membership queries to cache");
            equivalenceAlgorithm = new ModifiedWMethodEQOracle.MealyModifiedWMethodEQOracle<String, String>(config.max_depth, statsCachedMemOracle);
            //AcexLocalSuffixFinder suffixFinder = new AcexLocalSuffixFinder(AcexAnalyzers.BINARY_SEARCH_BWD, true, "Analyzer");
            //learningAlgorithm = new TTTLearnerMealy<String, String>(alphabet, cachedMemOracle,  AcexAnalyzers.LINEAR_FWD);
            learningAlgorithm = new ExtensibleLStarMealyBuilder<String, String>().withAlphabet(alphabet).withOracle(statsCachedMemOracle).create();
            log.logPhase("Using learning algorithm " + learningAlgorithm.getClass().getSimpleName());
            log.logPhase("Using equivalence algorithm " + equivalenceAlgorithm.getClass().getSimpleName());

            log.logPhase("Starting learning");

            SimpleProfiler.start("Total time");

            boolean learning = true;
            Counter round = new Counter("Rounds", "");
            round.increment();
            System.out.println("Starting round " + round.getCount());
            SimpleProfiler.start("Learning");
            learningAlgorithm.startLearning();
            SimpleProfiler.stop("Learning");

            MealyMachine<?, String, ?, String> hypothesis = learningAlgorithm.getHypothesisModel();
            DefaultQuery<String, Word<String>> tmpCE = null;
            while (learning) {
                // Write outputs
                writeDotModel(hypothesis, alphabet, config.output_dir + "/hypothesis_" + progress.getCount() +"_"+ round.getCount() + ".dot");

                // Search counter-example
                System.out.println("Searching for counter-example");
                SimpleProfiler.start("Searching for counter-example");
                DefaultQuery<String, Word<String>> counterExample = null;
                for (DefaultQuery<String, Word<String>> CE: CEList)
                {
                    if(!Objects.equals(CE.getOutput(),hypothesis.computeOutput(CE.getInput())))
                    {
                        counterExample = CE;
                        break;
                    }
                }
                if(counterExample==null)
                    counterExample = equivalenceAlgorithm.findCounterExample(hypothesis, alphabet);
                SimpleProfiler.stop("Searching for counter-example");

                if (counterExample == null) {
                    // No counter-example found, so done learning
                    learning = false;

                    // Write outputs
                    writeDotModel(hypothesis, alphabet, config.output_dir + "/learnedModel.dot");
                    writeDotModel(hypothesis, alphabet, config.output_dir + "/learnedModel_"+progress.getCount()+".dot");
                    //writeAutModel(hypothesis, alphabet, config.output_dir + "/learnedModel.aut");
                } else {
                    if (tmpCE == null)
                        tmpCE = counterExample;
                    else {
                        if (tmpCE.toString().equals(counterExample.toString())) {
                            // No counter-example found, so done learning
                            learning = false;

                            // Write outputs
                            writeDotModel(hypothesis, alphabet, config.output_dir + "/learnedModel.dot");
                            writeDotModel(hypothesis, alphabet, config.output_dir + "/learnedModel_"+progress.getCount()+".dot");
                            //writeAutModel(hypothesis, alphabet, config.output_dir + "/learnedModel.aut");
                        } else
                            tmpCE = counterExample;
                    }
                    CEList.add(counterExample);
                    // Counter example found, update hypothesis and continue learning
                    System.out.println("Counter-example found: " + counterExample.toString());
                    //TODO Add more logging
                    round.increment();
                    System.out.println("Starting round " + round.getCount());

                    SimpleProfiler.start("Learning");
                    learningAlgorithm.refineHypothesis(counterExample);
                    SimpleProfiler.stop("Learning");

                    hypothesis = learningAlgorithm.getHypothesisModel();
                }

                //System.out.println( "States in final hypothesis: " + hypothesis.size());

            }

            //deduplication
            MealyDot mealyAutoamta = new MealyDot(config.output_dir + "/learnedModel.dot");
            if(addtypes!=null)
            {
                Iterator iter = addtypes.entrySet().iterator();
                while (iter.hasNext()) {
                    Map.Entry entry = (Map.Entry) iter.next();
                    String key = (String)entry.getKey();
                    List<Integer> val = (List<Integer>)entry.getValue();
                    for(int index = val.size()-1; index > -1; index--)
                    {
                        String tmp = key+(new DecimalFormat("00").format(val.get(index)));
                        mealyAutoamta = Dedup(tmp, mealyAutoamta);
                    }
                }
            }
            //check the consistence of the flow and mealy state machine
            MealyState initialState = mealyAutoamta.getInitialState();
            addtypes = new LinkedHashMap<>();
            for(;i < inputTypeSeqs.size();i++)
            {
                List<String> inputList = inputTypeSeqs.get(i);
                List<String> outputList = outputTypeSeqs.get(i);
                ArrayList<List<MealyState>> stateArray = new ArrayList<>();
                List<MealyState> initList = new ArrayList<>();
                initList.add(initialState);
                stateArray.add(initList);
                Boolean isEqual = true;
                for(int j =0; j < inputList.size(); j++)
                {
                    List<MealyState> curList = new ArrayList<>();
                    for(MealyState curState: stateArray.get(j)){
                        String type = inputList.get(j);
                        if(!subtypes.containsKey(type))
                        {
                            isEqual = false;
                            break;
                        }
                        int size = subtypes.get(type);
                        for(int k = 0; k < size; k++)
                        {
                            String symbol = type + (new DecimalFormat("00").format(k+1));
                            OutputTransition outputTransition = curState.getOutputTransition(symbol);
                            if(outputTransition.output.equals(outputList.get(j)))
                            {
                                curList.add(outputTransition.state);
                            }
                        }
                    }
                    if(curList.isEmpty())
                    {
                        isEqual = false;
                        break;
                    }
                    stateArray.add(curList);
                    if(!isEqual)
                        break;
                }
                if(!isEqual) {
                    //construct alphabet and mapper
                    List<byte[]> seed = seedPool.get(i);
                    for(int j =0; j < inputList.size(); j++)
                    {
                        String type = inputList.get(j);
                        if(sul.addMapper(type,seed.get(j))) {
                            subtypes = sul.getSubtypes();
                            int size = subtypes.get(type);
                            List<Integer> tmp;
                            if (addtypes.containsKey(type)) {
                                tmp = addtypes.get(type);
                            } else {
                                tmp = new ArrayList<>();
                            }
                            tmp.add(size);
                            addtypes.put(type, tmp);
                        }
                    }
                    i++;
                    if(!addtypes.isEmpty())
                    {
                        alphabet = sul.getAlphabet();
                        break;
                    }
                }
            }
            if(i == inputTypeSeqs.size() && addtypes.isEmpty())
                break;

        }
        SimpleProfiler.stop("Total time");
        sul.saveMapper();
        MealyDot finalAutoamta = new MealyDot(config.output_dir + "/learnedModel.dot");
        subtypes = sul.getSubtypes();
        Iterator iter = subtypes.entrySet().iterator();
        while (iter.hasNext()) {
            Map.Entry entry = (Map.Entry) iter.next();
            String type = (String)entry.getKey();
            Integer val = (Integer)entry.getValue();
            for(int index = val; index > 1; index--)
            {
                finalAutoamta = new MealyDot(config.output_dir + "/learnedModel.dot");
                String tmp = type+(new DecimalFormat("00").format(index));
                finalAutoamta = DelSubtype(tmp, finalAutoamta);
                finalAutoamta.writeFile(config.output_dir + "/learnedModel.dot");
            }
        }
        // Output statistics
        System.out.println("-------------------------------------------------------");
        System.out.println( SimpleProfiler.getResults());
        System.out.println(progress.getSummary());
        System.out.println(statsCachedMemOracle.getStatisticalData().getSummary());
        System.out.println(sul.getTotalNum());
    }

    public static void writeDotModel(MealyMachine<?, String, ?, String> model, Alphabet<String> alphabet, String filename) throws IOException, InterruptedException {
        // Write output to dot-file
        File dotFile = new File(filename);
        PrintStream psDotFile = new PrintStream(dotFile);
        GraphDOT.write(model, alphabet, psDotFile);
        psDotFile.close();

        //TODO Check if dot is available

        // Convert .dot to .pdf
        Runtime.getRuntime().exec("dot -Tpdf -O " + filename);
    }

    public void initMapper() throws IOException, InterruptedException {
        inputTypeSeqs = new ArrayList<>();
        outputTypeSeqs = new ArrayList<>();
        seedPool = new ArrayList<>();
        seedPool = readFromFile("./Seed");
        List<List<byte[]>> tmpSeed = new ArrayList<>();
        tmpSeed.addAll(seedPool);
        int maxIndex=0;
        int maxLength = 0;
        try{
            for (List<byte[]> seed :tmpSeed) {
                sul.pre();
                boolean flag = false;
                int i = 0;
                List<String> inputList = new ArrayList<>();
                List<String> outputList = new ArrayList<>();
                List<String> temp = null;
                for (byte[] message : seed){
                    inputList.add(sul.messageToSymbol(message));
                    outputList.add(sul.sendFuzzingMessage(message));
                }
                for(i = 0; i < inputTypeSeqs.size(); i++)
                {
                    temp = inputTypeSeqs.get(i);
                    flag= true;
                    for(int j = 0 ; j < Math.min(temp.size(),inputList.size()); j++)
                    {
                        if(temp.get(j).equals(inputList.get(j)))
                            continue;
                        else
                        {
                            flag= false;
                            break;
                        }
                    }
                    //check the sequences of the output type
                    if(flag)
                    {
                        temp = outputTypeSeqs.get(i);
                        for(int j = 0 ; j < Math.min(temp.size(),outputList.size()); j++)
                        {
                            if(temp.get(j).equals(outputList.get(j)))
                                continue;
                            else
                            {
                                flag= false;
                                break;
                            }
                        }
                    }
                    if(flag)
                        break;
                }
                if(!flag)
                {
                    inputTypeSeqs.add(inputList);
                    outputTypeSeqs.add(outputList);
                    if(inputList.size()>maxLength)
                    {
                        maxIndex = inputTypeSeqs.size()-1;
                        maxLength = inputList.size();
                    }
                }
                else if(flag && temp.size() < outputList.size())
                {
                    inputTypeSeqs.remove(i);
                    outputTypeSeqs.remove(i);
                    inputTypeSeqs.add(inputList);
                    outputTypeSeqs.add(outputList);
                    seedPool.remove(i);
                    if(inputList.size()>maxLength)
                    {
                        maxIndex = inputTypeSeqs.size()-1;
                        maxLength = inputList.size();
                    }
                }
                else
                {
                    seedPool.remove(seed);
                }
            }
            //select the most long inputlist and construct alphabet
            Collections.swap(inputTypeSeqs,maxIndex,0);
            Collections.swap(seedPool,maxIndex,0);
            Collections.swap(outputTypeSeqs,maxIndex,0);
            List<String> typeList = inputTypeSeqs.get(0);
            List<byte[]> messages = seedPool.get(0);
            int k = 0;
            for(String symbol: typeList)
            {
                sul.addMapper(symbol,messages.get(k++));
                subtypes = sul.getSubtypes();
            }
            alphabet =sul.getAlphabet();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void deterMutate() throws Exception
    {
        int size = seedPool.size();
        int i = 0;
        int j = 0;
        int mutateNum = 0;
        System.out.println("State fuzzing");
        LinkedHashMap<Integer, List<byte[]>> standardValue = new LinkedHashMap<>();
        for(int stateIndex = 0; stateIndex < size; stateIndex++) {
            int exec_num = 0;
            int m = 0;
            byte[] message;
            System.out.println(stateIndex);
            String outputStr = outputTypeSeqs.get(stateIndex).toString();
            for (int mutateIndex = 0; mutateIndex < seedPool.get(stateIndex).size(); mutateIndex++) {
                List<byte[]> currentSeed = deepCopy(seedPool.get(stateIndex));
                List<byte[]> interesting;
                int start = sul.getIndexOfSymbol(currentSeed.get(mutateIndex));
                if (!standardValue.containsKey(byteToHex(currentSeed.get(mutateIndex)).hashCode())) {
                    interesting = new ArrayList<>();
                    for (j = start; j < currentSeed.get(mutateIndex).length-2; j++) {
                        sul.pre();
                        List<byte[]> mutateSeed = new ArrayList<>();
                        m = 0;
                        List<String> output = new ArrayList<>();
                        for (; m < currentSeed.size(); m++) {
                            if (m == mutateIndex) {
                                //mutate the mth package
                                message = new byte[currentSeed.get(m).length];
                                System.arraycopy(currentSeed.get(m), 0, message, 0, currentSeed.get(m).length);
                                message[j] ^= 0xFF;
                                output.add(sul.sendFuzzingMessage(message));
                                mutateSeed.add(message);
                            } else {
                                output.add(sul.sendFuzzingMessage(currentSeed.get(m)));
                                mutateSeed.add(currentSeed.get(m));
                            }
                            if (output.toString().contains("Closed"))
                            {
                                m++;
                                break;
                            }
                        }
                        sul.post();
                        if (!output.toString().equals(outputStr)) {
                            seedPool.add(mutateSeed);
                            inputTypeSeqs.add(inputTypeSeqs.get(stateIndex).subList(0,m));
                            outputTypeSeqs.add(output);
                            interesting.add(mutateSeed.get(Math.min(mutateIndex, mutateSeed.size() - 1)));
                            mutateNum ++;
                            seedSave(mutateSeed,   config.seed_dir + File.separator + mutateNum);
                        }
                    }
                } else {
                    interesting = standardValue.get(byteToHex(currentSeed.get(mutateIndex)).hashCode());
                    for (j = start; j < interesting.size(); j++) {
                        sul.pre();
                        List<byte[]> mutateSeed = new ArrayList<>();
                        m = 0;
                        List<String> output = new ArrayList<>();
                        for (; m < currentSeed.size(); m++) {
                            if (m == mutateIndex) {
                                //mutate the mth package
                                message = new byte[interesting.get(j).length];
                                System.arraycopy(interesting.get(j), 0, message, 0, interesting.get(j).length);
                                output.add(sul.sendFuzzingMessage(message));
                                mutateSeed.add(message);
                            } else {
                                output.add(sul.sendFuzzingMessage(currentSeed.get(m)));
                                mutateSeed.add(currentSeed.get(m));
                            }
                            if (output.toString().contains("Closed"))
                            {
                                m++;
                                break;
                            }
                        }
                        sul.post();
                        if (!output.toString().equals(outputStr)) {
                            seedPool.add(mutateSeed);
                            inputTypeSeqs.add(inputTypeSeqs.get(stateIndex).subList(0,m));
                            outputTypeSeqs.add(output);
                            interesting.add(mutateSeed.get(Math.min(mutateIndex, mutateSeed.size() - 1)));
                            mutateNum ++;
                            seedSave(mutateSeed,   config.seed_dir + File.separator + mutateNum);
                        }
                    }
                }
            }
        }
    }

    public void seedSave(List<byte[]> seed, String filename) {
        BufferedOutputStream bos = null;
        FileOutputStream fos = null;
        File file = null;
        try {
            file = new File(filename);
            if (file.exists()) {
                return;
            }
            fos = new FileOutputStream(file);
            bos = new BufferedOutputStream(fos);
            for (int i = 0; i < seed.size(); i++) {
                byte[] tmp = intToByteArray(seed.get(i).length);
                bos.write(tmp, 0, 4);
                bos.write(seed.get(i), 0, seed.get(i).length);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (bos != null) {
                try {
                    bos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
            if (fos != null) {
                try {
                    fos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
        }
    }

    public MealyDot Dedup(String subtype, MealyDot mealyAutoamta) throws IOException, InterruptedException
    {
        String type = subtype.substring(0,subtype.length()-2);
        MealyState initState = mealyAutoamta.getInitialState();
        Set<MealyState> states = mealyAutoamta.getStates();
        List<String> typeList = new ArrayList<>();
        for(int i = 1; i < subtypes.get(type); i++)
        {
            typeList.add(type+(new DecimalFormat("00").format(i)));
        }
        for(String tmp:typeList){
            boolean flag = false;
            for(MealyState state: states)
            {
                if(!mealyAutoamta.getNextState(state,tmp).getName().equals(mealyAutoamta.getNextState(state,subtype).getName()) ||
                        !mealyAutoamta.getOutput(state,tmp).equals(mealyAutoamta.getOutput(state,subtype)))
                {
                    flag = true;
                    break;
                }
            }
            if(!flag)
            {
                sul.delMapper(subtype);
                //update the mealy
                mealyAutoamta.updateMealy(subtype);
                logEqOracle.updateLog(subtype);
                dedupFlag = true;
                break;
            }
        }
        return mealyAutoamta;
    }

    public MealyDot DelSubtype(String subtype, MealyDot mealyAutoamta) throws IOException, InterruptedException
    {
        String type = subtype.substring(0,subtype.length()-2);
        mealyAutoamta.writeFile("tmp.dot");
        MealyDot tmpMealy = new MealyDot("tmp.dot");
        MealyState initState = tmpMealy.getInitialState();
        Set<MealyState> states = tmpMealy.getStates();
        tmpMealy.updateMealy(subtype);
        if(tmpMealy.getStatesNum() < mealyAutoamta.getStatesNum())
            return mealyAutoamta;
        else
        {
            sul.delMapper(subtype);
            return tmpMealy;
        }
    }

    public void test() throws Exception
    {
        MealyDot mealyAutoamta = new MealyDot("learnedModel.dot");
        Set<String> tmp = mealyAutoamta.getInputAlphabet();
        List<String> tmpList = new ArrayList<>(tmp);
        for(int i = tmpList.size()-1; i > -1;i--){
            mealyAutoamta = new MealyDot("learnedModel.dot");
            String symbol = tmpList.get(i);
            mealyAutoamta = DelSubtype(symbol, mealyAutoamta);
            mealyAutoamta.writeFile("learnedModel.dot");
        }
        mealyAutoamta.writeFile("learnedModel.dot");
        System.out.println("Done");
    }


    public static void main(String[] args) throws Exception {
        if(args.length < 1) {
            System.err.println("Invalid number of parameters");
            System.exit(-1);
        }
        LearningConfig config = new LearningConfig(args[0]);

        ProLearner learner = new ProLearner(config);
        learner.learn();

        System.exit(0);
    }
}
